#!/usr/bin/env python3
# Copyright 2024 The Emscripten Authors.  All rights reserved.
# Emscripten is available under two separate licenses, the MIT license and the
# University of Illinois/NCSA Open Source License.  Both these licenses can be
# found in the LICENSE file.

"""Automatically rebaseline tests that have codesize expectations and create
a git commit containing the resulting changes along with readable details of
the generated changes.
"""

import argparse
import json
import os
import statistics
import subprocess
import sys

script_dir = os.path.dirname(os.path.abspath(__file__))
root_dir = os.path.dirname(os.path.dirname(script_dir))

sys.path.insert(0, root_dir)
from tools import utils


def run(cmd, **args):
  return subprocess.check_output(cmd, text=True, cwd=root_dir, **args)


all_deltas = []


def read_size_from_json(content):
  json_data = json.loads(content)
  if 'total' in json_data:
    return json_data['total']
  # If `total` if not in the json dict then just use the first key.  This happens when only one
  # file size is reported (in this case we don't calculate or store the `total`).
  first_key = list(json_data.keys())[0]
  return json_data[first_key]


def process_changed_file(filename):
  content = open(filename).read()
  old_content = run(['git', 'show', f'HEAD:{filename}'])
  print(f'processing {filename}')

  ext = os.path.splitext(filename)[1]
  if ext == '.size':
    size = int(content.strip())
    old_size = int(old_content.strip())
  elif ext == '.json':
    size = read_size_from_json(content)
    old_size = read_size_from_json(old_content)
  else:
    # Unhandled file type
    return f'{filename} updated\n'

  filename = filename.removeprefix('test/')
  delta = size - old_size
  percent_delta = delta * 100 / old_size
  all_deltas.append(percent_delta)
  return f'{filename}: {old_size} => {size} [{delta:+} bytes / {percent_delta:+.2f}%]\n'


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('-s', '--skip-tests', action='store_true', help="Don't actually run the tests, just analyze the existing results")
  parser.add_argument('-b', '--new-branch', action='store_true', help='Create a new branch containing the updates')
  parser.add_argument('-c', '--clear-cache', action='store_true', help='Clear the cache before rebaselining (useful when working with llvm changes)')
  parser.add_argument('-n', '--check-only', dest='check_only', action='store_true', help='Return non-zero if test expectations are out of date, and skip creating a git commit')
  args = parser.parse_args()

  if args.clear_cache:
    run(['./emcc', '--clear-cache'])

  if not args.skip_tests:
    if not args.check_only and run(['git', 'status', '-uno', '--porcelain']).strip():
      print('tree is not clean')
      return 1

    subprocess.check_call([utils.bat_suffix(os.path.join('test', 'runner')), '--rebaseline', 'codesize'], cwd=root_dir)

  output = run(['git', 'status', '-uno', '--porcelain'])
  filenames = []
  for line in output.splitlines():
    filename = line.strip().rsplit(' ', 1)[1]
    if filename.startswith('test') and os.path.isfile(filename):
      filenames.append(filename)

  if not filenames:
    print('test expectations are up-to-date')
    return 0

  if args.check_only:
    message = f'''Test expectations are out-of-date

The following ({len(filenames)}) test expectation files were updated by
running the tests with `--rebaseline`:

```
'''
  else:
    message = f'''
Automatic rebaseline of codesize expectations. NFC

This is an automatic change generated by tools/maint/rebaseline_tests.py.

The following ({len(filenames)}) test expectation files were updated by
running the tests with `--rebaseline`:

```
'''

  for file in filenames:
    message += process_changed_file(file)

  if all_deltas:
    message += f'\nAverage change: {statistics.mean(all_deltas):+.2f}% ({min(all_deltas):+.2f}% - {max(all_deltas):+.2f}%)\n'

  message += '```\n'

  print(message)
  if args.check_only:
    return 1

  if args.new_branch:
    run(['git', 'checkout', '-b', 'rebaseline_tests'])
  run(['git', 'add', '-u', '.'])
  run(['git', 'commit', '-F', '-'], input=message)

  return 2


if __name__ == '__main__':
  sys.exit(main())
